// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/api.h>
#include <nebula/compute/api.h>
#include <nebula/csv/api.h>
#include <nebula/csv/writer.h>
#include <nebula/io/api.h>

#include <turbo/utility/status.h>

#include <iostream>
#include <vector>

// Many operations in Apache Nebula operate on
// columns of data, and the columns of data are
// assembled into a table. In this example, we
// examine how to compare two arrays which are
// combined to form a table that is then written
// out to a CSV file.
//
// To run this example you can use
// ./compute_and_write_csv_example
//
// the program will write the files into
// compute_and_write_output.csv
// in the current directory

turbo::Status RunMain(int argc, char **argv) {
    // Make Arrays
    nebula::NumericBuilder<nebula::Int64Type> int64_builder;
    nebula::NumericBuilder<nebula::Int64Type> sum_builder;

    // Make place for 8 values in total
    TURBO_RETURN_NOT_OK(int64_builder.resize(8));
    TURBO_RETURN_NOT_OK(sum_builder.resize(8));

    // Bulk append the given values
    std::vector<int64_t> int64_values = {1, 2, 3, 4, 5, 6, 7, 8};
    TURBO_RETURN_NOT_OK(int64_builder.append_values(int64_values));
    std::shared_ptr<nebula::Array> array_a;
    TURBO_RETURN_NOT_OK(int64_builder.finish(&array_a));
    int64_builder.reset();
    int64_values = {2, 5, 1, 3, 6, 2, 7, 4};
    std::shared_ptr<nebula::Array> array_b;
    TURBO_RETURN_NOT_OK(int64_builder.append_values(int64_values));
    TURBO_RETURN_NOT_OK(int64_builder.finish(&array_b));

    // Cast the arrays to their actual types
    auto int64_array_a = std::static_pointer_cast<nebula::Int64Array>(array_a);
    auto int64_array_b = std::static_pointer_cast<nebula::Int64Array>(array_b);
    // Explicit comparison of values using a loop
    for (int64_t i = 0; i < 8; i++) {
        if ((!int64_array_a->is_null(i)) && (!int64_array_b->is_null(i))) {
            auto comparison_result = int64_array_a->value(i) + int64_array_b->value(i);
            sum_builder.unsafe_append(comparison_result);
        } else {
            sum_builder.unsafe_append_null();
        }
    }
    std::shared_ptr<nebula::Array> array_a_gt_b_self;
    TURBO_RETURN_NOT_OK(sum_builder.finish(&array_a_gt_b_self));
    std::cout << "Array explicitly compared" << std::endl;

    // Explicit comparison of values using a compute function
    TURBO_MOVE_OR_RAISE(nebula::Datum compared_datum,
            nebula::compute::call_function("add", {array_a, array_b}));
    auto array_a_gt_b_compute = compared_datum.make_array();
    std::cout << "Arrays compared using a compute function" << std::endl;

    // create a table for the output
    auto schema =
            nebula::schema({nebula::field("a", nebula::int64()), nebula::field("b", nebula::int64()),
                            nebula::field("a+b (self written)", nebula::int64()),
                            nebula::field("a+b (nebula)", nebula::int64())});
    std::shared_ptr<nebula::Table> my_table = nebula::Table::create(
            schema, {array_a, array_b, array_a_gt_b_self, array_a_gt_b_compute});

    std::cout << "Table created" << std::endl;

    // write table to CSV file
    auto csv_filename = "compute_add_output.csv";
    TURBO_MOVE_OR_RAISE(auto outstream, nebula::io::FileOutputStream::open(csv_filename));

    std::cout << "Writing CSV file" << std::endl;
    TURBO_RETURN_NOT_OK(nebula::csv::write_csv(
            *my_table, nebula::csv::WriteOptions::defaults(), outstream.get()));

    return turbo::OkStatus();
}

int main(int argc, char **argv) {
    turbo::Status status = RunMain(argc, argv);
    if (!status.ok()) {
        std::cerr << status.to_string() << std::endl;
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}
